import datetime
import requests

from map import map

class BaiduIndex:
    """
    使用该类完成百度指数的获取与存储
    """
    def __init__(self):
        """
        类的初始化，完成一些基本参数的配置
        """
        self.base_url = "http://index.baidu.com/api/SearchApi/index"
        self.headers = {
            'Accept': 'application/json, text/plain, */*',
            'Accept-Encoding': 'gzip, deflate',
            'Accept-Language': 'zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6',
            'Cache-Control': 'no-cache',
            'Connection': 'keep-alive',
            'Host': 'index.baidu.com',
            'Pragma': 'no-cache',
            'Referer': 'http://index.baidu.com/v2/main/index.html',
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/80.0.3987.132 Safari/537.36 Edg/80.0.361.66',
            'X-Requested-With': 'XMLHttpRequest',
        }
        self.session = requests.Session()

        self.map = map

    def login(self, cookies: str) -> None:
        """
        使用用户提供的cookies模拟登录
        """
        self.headers.update({
            'Cookie': 'BDUSS=' + cookies
        })

    def search(self, keyword: str, area: str, startDate: str, endDate: str, dataType: str = 'all', multi_keys: bool = False) -> list:
        """
        获取具体地区某个关键词的百度指数
        @param keyword: 待查询的关键词，如果需要同时查询多个关键词请使用 + 将其分隔开，比如 狗不理+王府井+门店
        @param area: 待查询的区域名，具体可用的名字请查阅当前文件夹下的map.py
        @param startDate: 开始日期
        @param endDate: 结束日期 开始日期与结束日期之间的间隔不要超过一年，不然返回的内容是以周为周期的指数而非日度指数了
        @param dataType: 数据类型，取值为all时代表PC加移动，pc代表PC端，wise代表移动端
        @param multi_keys: 是否同时查询多个关键词，默认取值为False，表明为只查询单个关键词，如果需要同时查询多个关键词，需要将该参数修改为True

        @return list: 解密后的数据列表
        """
        self.startDate = startDate
        
        if multi_keys:
            word_list = keyword.split('+')
            word = [[{"name": w, "wordType": 1} for w in word_list]]
        else:
            word = [[{"name":  keyword, "wordType": 1}]]
        params = {
            'area': self.map[area],
            'word': str(word).replace("'", '"'),
            'startDate': startDate,
            'endDate': endDate,
        }

        resp = self.session.get(url=self.base_url, headers=self.headers, params=params).json()

        passwd = self._get_passwd(resp['data']['uniqid'])
        data = resp['data']['userIndexes'][0][dataType]['data']
        return self.decrypt(passwd, data)

    def _get_passwd(self, uniqid: str) -> str:
        """
        获取解密密文的密码
        @param uniqid 获取密码的必要参数

        @return str 从服务器处取回的密码
        """
        url = f'http://index.baidu.com/Interface/ptbk?uniqid={uniqid}'
        return self.session.get(url=url, headers=self.headers).json()['data']

    @staticmethod
    def decrypt(password, data) -> list:
        """
        解密图表数据
        @param password 密码表数据
        @param data 待解密数据
        @return 解密后的数据
        """
        arr = list(password)
        dataArr = list(data)
        objPass = {}
        notInNames = []
        for i in range(len(arr)//2):
            objPass[arr[i]] = arr[len(arr) // 2 + i]

        # 数据解密转换
        for i in range(len(dataArr)):
            notInNames.append(objPass[dataArr[i]])

        str_list = "".join(notInNames).split(",")
        return [int(s) if s != '' else 0 for s in str_list]

    def save(self, data: list, file_path: str) -> None:
        """
        将获取到的百度指数存储在指定位置
        """
        # 将列表数据与起始日期串起来
        start_date = datetime.datetime.strptime(self.startDate, r'%Y-%m-%d')
        index = 0
        result = ""
        for item in data:
            result += (start_date + datetime.timedelta(days=index)).strftime(r'%Y-%m-%d') + ',' + str(item) + '\n'
            index += 1

        with open(file_path, 'w') as f:
            f.write(result)
        
        print("Job done!")

if __name__ == "__main__":
    app = BaiduIndex()
    app.login(r'')
    data = app.search("狗不理+王府井+门店", '南昌', '2019-1-1', '2019-12-31', 'all', multi_keys=True)
    app.save(data, 'nanchang.csv')
